import argparse
import logging
import os
import random
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
from pathlib import Path
from torch import optim
from torch.utils.data import DataLoader, random_split
from tqdm import tqdm

import wandb

from data_loading import EMDataset, BasicDataset
from evaluate import evaluate
from img_util import get_points_image, get_mask_image
from loss import dist_loss, balance_loss, label_accuracy_loss, get_loss
from unet_model_simple import UNet

from util import predict_mask

# dir_img = Path('/LSEM/user/zhoucan/unet/train/')
dir_img = Path('./data/imgs/')
# dir_mask = Path('/LSEM/user/zhoucan/unet/train_masks/')
dir_mask = Path('./data/masks/')
# dir_checkpoint = Path('/LSEM/user/zhoucan/unet/checkpoints')
dir_checkpoint = Path('./checkpoints/')


def train_model(
        model,
        device,
        dataset,
        n_train,
        n_val,
        train_loader,
        val_loader,
        epochs: int = 5,
        batch_size: int = 1,
        learning_rate: float = 1e-5,
        val_percent: float = 0.1,
        save_checkpoint: bool = True,
        img_scale: float = 0.5,
        amp: bool = False,
        weight_decay: float = 1e-8,
        momentum: float = 0.999,
        gradient_clipping: float = 1.0,
):
    # (Initialize logging)
    experiment = wandb.init(project='EM_point_hinter_local', resume='allow', anonymous='must')
    experiment.config.update(
        dict(epochs=epochs, batch_size=batch_size, learning_rate=learning_rate,
             val_percent=val_percent, save_checkpoint=save_checkpoint, img_scale=img_scale, amp=amp)
    )

    logging.info(f'''Starting training:
        Epochs:          {epochs}
        Batch size:      {batch_size}
        Learning rate:   {learning_rate}
        Training size:   {n_train}
        Validation size: {n_val}
        Checkpoints:     {save_checkpoint}
        Device:          {device.type}
        Images scaling:  {img_scale}
        Mixed Precision: {amp}
    ''')

    # 4. Set up the optimizer, the loss, the learning rate scheduler and the loss scaling for AMP
    optimizer = optim.RMSprop(model.parameters(),
                              lr=learning_rate, weight_decay=weight_decay, momentum=momentum, foreach=True)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=5)  # goal: minimize val loss
    grad_scaler = torch.cuda.amp.GradScaler(enabled=amp)
    # criterion = nn.CrossEntropyLoss() if model.n_classes > 1 else nn.BCEWithLogitsLoss()
    global_step = 0
    # 5. Begin training
    for epoch in range(1, epochs + 1):
        model.train()
        epoch_loss = 0
        with tqdm(total=n_train, desc=f'Epoch {epoch}/{epochs}', unit='img') as pbar:
            for batch in train_loader:
                images, true_masks = batch['image'], batch['mask']

                assert images.shape[1] == model.n_channels, \
                    f'Network has been defined with {model.n_channels} input channels, ' \
                    f'but loaded images have {images.shape[1]} channels. Please check that ' \
                    'the images are loaded correctly.'

                images = images.to(device=device, dtype=torch.float32, memory_format=torch.channels_last)
                true_masks = true_masks.to(device=device, dtype=torch.long)

                with torch.autocast(device.type if device.type != 'mps' else 'cpu', enabled=amp):
                    masks_pred = model(images)
                    # masks_pred = model(images).squeeze(1)
                    logging.info('masks_pred: {}'.format(torch.sigmoid(masks_pred)))
                    dist_loss, balance_loss, label_loss = get_loss(masks_pred, true_masks)
                    loss = dist_loss + balance_loss + label_loss

                optimizer.zero_grad(set_to_none=True)
                grad_scaler.scale(loss).backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clipping)
                grad_scaler.step(optimizer)
                grad_scaler.update()

                pbar.update(images.shape[0])
                global_step += 1
                epoch_loss += loss.item()
                experiment.log({
                    'train loss': loss.item(),
                    'label_loss': label_loss,
                    'balance_loss': balance_loss,
                    'dist_loss': dist_loss,
                    'step': global_step,
                    'epoch': epoch
                })
                pbar.set_postfix(**{'loss (batch)': loss.item()})

                # Evaluation round
                division_step = (n_train // (200 * batch_size))
                if division_step > 0:
                    if global_step % division_step == 0:
                        histograms = {}
                        for tag, value in model.named_parameters():
                            tag = tag.replace('/', '.')
                            if not (torch.isinf(value) | torch.isnan(value)).any():
                                histograms['Weights/' + tag] = wandb.Histogram(value.data.cpu())
                            if not (torch.isinf(value.grad) | torch.isnan(value.grad)).any():
                                histograms['Gradients/' + tag] = wandb.Histogram(value.grad.data.cpu())

                        val_loss = evaluate(model, val_loader, device, amp)
                        scheduler.step(val_loss)

                        logging.info('Validation loss: {}'.format(val_loss))
                        # try:
                        # print("true_masks_shape", true_masks.shape)
                        # print("masks_pred_shape", masks_pred.shape)
                        # print("masks_pred.argmax(dim=1)shape", masks_pred.argmax(dim=1).shape)
                        # print("masks_pred", masks_pred)
                        # print("true_masks", true_masks[0][500])
                        # print("true_masks", true_masks[0][600])
                        # print("true_masks", true_masks[0][400])
                        result = masks_pred[0].detach().cpu()

                        logging.info('result: {}'.format(result))
                        result = torch.sigmoid(result)

                        logging.info('points_pred:{}'.format(result))

                        # fake_result = torch.tensor([[0.5126, 0.5126, 0.6315],
                        # [0.9, 0.9, 0.3]])

                        # result = fake_result
                        mask, score = predict_mask(images[0], result)

                        logging.info('score:{}'.format(score))

                        image = get_mask_image(images[0].cpu(), mask, result)
                        experiment.log({
                            'learning rate': optimizer.param_groups[0]['lr'],
                            'validation Loss': val_loss,
                            'images': wandb.Image(images[0].cpu()),
                            'masks': {
                                'true': wandb.Image(true_masks[0].float().cpu()),
                                'pred': wandb.Image(get_points_image(images[0].cpu(), result)),
                                'pred_mask': wandb.Image(image),
                                # 'pred': wandb.Image(masks_pred.argmax(dim=1)[0].float().cpu()),
                            },
                            'step': global_step,
                            'epoch': epoch,
                            **histograms
                        })
                    # except:
                    #     pass

    if save_checkpoint:
        Path(dir_checkpoint).mkdir(parents=True, exist_ok=True)
        state_dict = model.state_dict()
        state_dict['mask_values'] = dataset.mask_values
        torch.save(state_dict, str(dir_checkpoint / 'checkpoint_epoch{}.pth'.format(epoch)))
        logging.info(f'Checkpoint {epoch} saved!')


# def apply_sigmoid_to_points(points):
#     # 假设points的形状是(1, 3)，其中前两个元素是坐标
#     coords = points[:, :2]  # 取出前两个坐标
#     sigmoid_coords = torch.sigmoid(coords)  # 对这些坐标应用Sigmoid函数
#     return torch.cat((sigmoid_coords, points[:, 2:]), dim=-1)  # 将处理后的坐标与第三个元素重新组合

def split_train_data(batch_size, img_scale, val_percent):
    # 1. Create dataset
    try:
        dataset = EMDataset(dir_img, dir_mask, img_scale)
    except (AssertionError, RuntimeError, IndexError):
        dataset = BasicDataset(dir_img, dir_mask, img_scale)
    # 2. Split into train / validation partitions
    n_val = int(len(dataset) * val_percent)
    n_train = len(dataset) - n_val
    train_set, val_set = random_split(dataset, [n_train, n_val], generator=torch.Generator().manual_seed(0))
    # 3. Create data loaders
    loader_args = dict(batch_size=batch_size, num_workers=os.cpu_count(), pin_memory=True)
    train_loader = DataLoader(train_set, shuffle=True, **loader_args)
    val_loader = DataLoader(val_set, shuffle=False, drop_last=True, **loader_args)
    return dataset, n_train, n_val, train_loader, val_loader


def initialize_model(model, train_loader, device="cpu"):
    """
    使用真实的批次数据初始化模型。

    参数:
        model: 要初始化的模型。
        train_loader: 用于提取数据的数据加载器。
        device: 运行模型的设备（例如 'cuda' 或 'cpu'）。
    """
    first_batch = next(iter(train_loader))
    images = first_batch['image'].to(device=device, dtype=torch.float32)
    _, height, width = images[0].shape
    with torch.no_grad():  # 确保不计算梯度
        model(images)


def get_args():
    parser = argparse.ArgumentParser(description='Train the UNet on images and target masks')
    parser.add_argument('--epochs', '-e', metavar='E', type=int, default=5, help='Number of epochs')
    parser.add_argument('--batch-size', '-b', dest='batch_size', metavar='B', type=int, default=1, help='Batch size')
    parser.add_argument('--learning-rate', '-l', metavar='LR', type=float, default=1e-5,
                        help='Learning rate', dest='lr')
    parser.add_argument('--load', '-f', type=str, default=False, help='Load model from a .pth file')
    parser.add_argument('--scale', '-s', type=float, default=0.5, help='Downscaling factor of the images')
    parser.add_argument('--validation', '-v', dest='val', type=float, default=10.0,
                        help='Percent of the data that is used as validation (0-100)')
    parser.add_argument('--amp', action='store_true', default=False, help='Use mixed precision')
    parser.add_argument('--bilinear', action='store_true', default=False, help='Use bilinear upsampling')
    # parser.add_argument('--classes', '-c', type=int, default=2, help='Number of classes')

    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()

    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logging.info(f'Using device {device}')

    dataset, n_train, n_val, train_loader, val_loader = split_train_data(args.batch_size, args.scale, args.val)
    # Change here to adapt to your data
    # n_channels=3 for RGB images
    # n_classes is the number of probabilities you want to get per pixel
    model = UNet(n_channels=3, bilinear=args.bilinear)
    initialize_model(model, train_loader)
    model = model.to(memory_format=torch.channels_last)

    logging.info(f'Network:\n'
                 f'\t{model.n_channels} input channels\n'
                 f'\t{"Bilinear" if model.bilinear else "Transposed conv"} upscaling')

    if args.load:
        state_dict = torch.load(args.load, map_location=device)
        del state_dict['mask_values']
        model.load_state_dict(state_dict)
        logging.info(f'Model loaded from {args.load}')

    model.to(device=device)
    # print(model.inc.conv[0].weight)
    torch.autograd.set_detect_anomaly(True)

    try:
        train_model(
            model=model,
            epochs=args.epochs,
            dataset=dataset, n_train=n_train, n_val=n_val, train_loader=train_loader, val_loader=val_loader,
            batch_size=args.batch_size,
            learning_rate=args.lr,
            device=device,
            img_scale=args.scale,
            val_percent=args.val / 100,
            amp=args.amp
        )
    except torch.cuda.OutOfMemoryError:
        logging.error('Detected OutOfMemoryError! '
                      'Enabling checkpointing to reduce memory usage, but this slows down training. '
                      'Consider enabling AMP (--amp) for fast and memory efficient training')
        torch.cuda.empty_cache()
        model.use_checkpointing()
        train_model(
            model=model,
            epochs=args.epochs,
            dataset=dataset, n_train=n_train, n_val=n_val, train_loader=train_loader, val_loader=val_loader,
            batch_size=args.batch_size,
            learning_rate=args.lr,
            device=device,
            img_scale=args.scale,
            val_percent=args.val / 100,
            amp=args.amp
        )
